from yolov8_video_detector.detector import VideoObjectDetector
from yolov8_video_detector.utils import (
    validate_video_file,
    get_video_properties,
    save_detections_to_json
)
import time


def main():
    input_video = "./input_video.mp4"
    output_video = "./output_video.mp4"
    model_path = "WSODD.pt"
    conf_threshold = 0.5
    show_video = True
    save_output = True
    output_fps = None
    json_output_path = "result.json"

    # 校验输入视频
    if not validate_video_file(input_video):
        print(f"Error: Could not open video file {input_video}")
        return

    # 打印视频信息
    video_info = get_video_properties(input_video)
    if video_info:
        print(f"Video Info:")
        print(f"  Resolution: {video_info['width']}x{video_info['height']}")
        print(f"  FPS: {video_info['fps']:.2f}")
        print(f"  Frames: {video_info['frame_count']}")
        print(f"  Duration: {video_info['duration']:.2f} seconds")

    # 初始化模型
    detector = VideoObjectDetector(model_path)
    print(f"Using model: {model_path}")
    print(f"Detectable classes: {', '.join(detector.get_class_names())}")

    start_time = time.time()

    # 目标检测
    success, message, video_detections = detector.detect_video(
        video_path=input_video,
        output_path=output_video,
        conf_threshold=conf_threshold,
        show=show_video,
        save=save_output,
        fps=output_fps
    )

    print(message)

    # 保存 JSON 结构化结果
    if success and json_output_path:
        json_saved = save_detections_to_json(video_detections, json_output_path)
        if json_saved:
            print(f"Detections saved to JSON: {json_output_path}")
        else:
            print("Failed to save detections to JSON.")

    end_time = time.time()
    print("Processing time: {:.2f} seconds".format(end_time - start_time))


if __name__ == '__main__':
    main()
